## DGP for Bayesian sensitivity analysis 
library(MASS)

simulate <- function(N = 200,
	                 TT = 50,
	                 Ntr = 50,
	                 T0 = 45,
	                 p = 2,
	                 beta = c(1, 3),
	                 r = 2,
	                 twoway = 1,
	                 c12 = 0.25,
	                 c22 = 0.25,
	                 c32 = 0.25,
	                 c42 = 0.25,
	                 c52 = 0.25,
	                 coeff = 0.5,
	                 betau = 1,
	                 ATT = 1) {

	## factor term 
	L <- matrix(rnorm(N*r), N, r)
	F <- matrix(rnorm(TT*r), TT, r)
	FL <- F %*% t(L)

	## two-way fe
	v.alpha <- rnorm(N)
	v.xi <- rnorm(TT)

	alpha <- as.matrix(rep(1, TT)) %*%  t(as.matrix(v.alpha))
	xi <- as.matrix(v.xi) %*% t(as.matrix(rep(1, N)))

	X <- array(NA, dim = c(TT, N, p))
	X.fit <- matrix(0, TT, N)
	for (i in 1:p) {
		X[,, i] <- matrix(rnorm(TT*N), TT, N) 
		X.fit <- X.fit + X[,,i] * beta[i]
	}

	D <- matrix(0, TT, N)
	D[(T0+1):TT, 1:Ntr] <- 1 

	## generate U 
    ## Sigma <- matrix(c(c22, -c32/2, -c32/2, c32), 2, 2)
	## l0 <- mvrnorm(n = 1, rep(0, 2), Sigma, tol = 1e-6, 
	## 	          empirical = FALSE, EISPACK = FALSE)
	## U <- l0[1] + l0[2] * D 
	U <- 0.3 - 0.5 * D  

	## 
	## lb <- rnorm(p, mean = 0, sd = sqrt(c42/p))
	lb <- rep(0.1, p)
	for (i in 1:p) {
		U <- U + X[,,i] * lb[i]
	}

	if (twoway) {
		## U <- U + rnorm(1, mean = 0, sd = sqrt(c52/(2*(r+2)))) * (FL + alpha + xi)
	    U <- U + 0.1 * (FL + alpha + xi)
	} else {
		## U <- U + rnorm(1, mean = 0, sd = sqrt(c52/(2*r))) * FL
	    U <- U + 0.1 * FL
	}

	## error 
	eu <- matrix(rnorm(TT*N, mean = 0, sd = sqrt(c12)), TT, N)
	U <- U + eu

	## error 
	e <- matrix(rnorm(TT*N, sd = 2), TT, N)

	## effect 
	eff <- matrix(0, TT, N)
	if (ATT != 0) {
		eff[(T0+1):TT,] <- matrix(rep(coeff * 1:(TT - T0), N), ,N)
	    eff <- eff * D 
    }
    veff <- apply(eff[,1:Ntr], 1, mean)


	Y <- e + FL + X.fit + eff + 3 + betau * U 
	if (twoway) {
		Y <- Y + alpha + xi
	}

	## data frame 
	data <- cbind.data.frame(rep(1:N, each = TT), rep(1:TT, N), 
		                     c(Y), c(D), c(U), c(eff), c(e))
	vx <- matrix(NA, TT * N, p)
	for (i in 1:p) {
		vx[, i] <- c(X[,, i])
	}
	data <- cbind.data.frame(data, vx)
	names(data) <- c("id", "time", "Y", "D", "U", "eff", "e", paste("X", 1:p, sep = ""))

	return(list(data = data, veff = veff))

}


## summarize treatment effect for simulation studies
effSummary_simple <- function(x) {  

    burn <- 0
    niter <- dim(x$sigma2)[2] 

    id.tr <- x$raw.id.tr
    time.tr <- x$time.tr
    rela.time.tr <- x$rela.time.tr  


    ## id indicator
    id.pos <- 1:length(c(id.tr))  
    

    yo_t <- x$yo_t
    yo_t <- yo_t[id.pos]

    time.tr <- time.tr[id.pos]
    rela.time.tr <- rela.time.tr[id.pos]

    yct_i <- x$yct
    yct_i <- matrix(c(yct_i[id.pos, (burn + 1):niter]), length(id.pos), niter - burn)


    ## average effetcs ------------

    t.post <- which(rela.time.tr > 0)

    eff_avg_i <- sapply(1:(niter - burn), function(i) {mean(yo_t[t.post] - yct_i[t.post, i])})

    out <- list(est.avg = eff_avg_i)

    return(out)

}